// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <memory>

#include <nebula/csv/converter.h>
#include <nebula/csv/options.h>
#include <turbo/log/logging.h>

namespace nebula::csv {

    enum class InferKind {
        Null,
        Integer,
        Boolean,
        Real,
        Date,
        Time,
        Timestamp,
        TimestampNS,
        TimestampWithZone,
        TimestampWithZoneNS,
        TextDict,
        BinaryDict,
        Text,
        Binary
    };

    class InferStatus {
    public:
        explicit InferStatus(const ConvertOptions &options)
                : kind_(InferKind::Null), can_loosen_type_(true), options_(options) {}

        InferKind kind() const { return kind_; }

        bool can_loosen_type() const { return can_loosen_type_; }

        void LoosenType(const turbo::Status &conversion_error) {
                    DKCHECK(can_loosen_type_);

            switch (kind_) {
                case InferKind::Null:
                    return SetKind(InferKind::Integer);
                case InferKind::Integer:
                    return SetKind(InferKind::Boolean);
                case InferKind::Boolean:
                    return SetKind(InferKind::Date);
                case InferKind::Date:
                    return SetKind(InferKind::Time);
                case InferKind::Time:
                    return SetKind(InferKind::Timestamp);
                case InferKind::Timestamp:
                    return SetKind(InferKind::TimestampNS);
                case InferKind::TimestampNS:
                    return SetKind(InferKind::TimestampWithZone);
                case InferKind::TimestampWithZone:
                    return SetKind(InferKind::TimestampWithZoneNS);
                case InferKind::TimestampWithZoneNS:
                    return SetKind(InferKind::Real);
                case InferKind::Real:
                    if (options_.auto_dict_encode) {
                        return SetKind(InferKind::TextDict);
                    } else {
                        return SetKind(InferKind::Text);
                    }
                case InferKind::TextDict:
                    if (turbo::is_out_of_range(conversion_error)) {
                        // Cardinality too large, fall back to non-dict encoding
                        return SetKind(InferKind::Text);
                    } else {
                        // Assuming UTF8 validation failure
                        return SetKind(InferKind::BinaryDict);
                    }
                    break;
                case InferKind::BinaryDict:
                    // Assuming cardinality too large
                    return SetKind(InferKind::Binary);
                case InferKind::Text:
                    // Assuming UTF8 validation failure
                    return SetKind(InferKind::Binary);
                default:
                    KLOG(FATAL) << "Shouldn't come here";
            }
        }

        turbo::Result<std::shared_ptr<Converter>> make_converter(MemoryPool *pool) {
            auto make_converter =
                    [&](std::shared_ptr<DataType> type) -> turbo::Result<std::shared_ptr<Converter>> {
                        return Converter::create(type, options_, pool);
                    };

            auto make_dict_converter =
                    [&](std::shared_ptr<DataType> type) -> turbo::Result<std::shared_ptr<Converter>> {
                        TURBO_MOVE_OR_RAISE(auto dict_converter,
                                               DictionaryConverter::create(type, options_, pool));
                        dict_converter->SetMaxCardinality(options_.auto_dict_max_cardinality);
                        return dict_converter;
                    };

            switch (kind_) {
                case InferKind::Null:
                    return make_converter(null());
                case InferKind::Integer:
                    return make_converter(int64());
                case InferKind::Boolean:
                    return make_converter(boolean());
                case InferKind::Date:
                    return make_converter(date32());
                case InferKind::Time:
                    return make_converter(time32(TimeUnit::SECOND));
                case InferKind::Timestamp:
                    return make_converter(timestamp(TimeUnit::SECOND));
                case InferKind::TimestampNS:
                    return make_converter(timestamp(TimeUnit::NANO));
                case InferKind::TimestampWithZone:
                    return make_converter(timestamp(TimeUnit::SECOND, "UTC"));
                case InferKind::TimestampWithZoneNS:
                    return make_converter(timestamp(TimeUnit::NANO, "UTC"));
                case InferKind::Real:
                    return make_converter(float64());
                case InferKind::Text:
                    return make_converter(utf8());
                case InferKind::Binary:
                    return make_converter(binary());
                case InferKind::TextDict:
                    return make_dict_converter(utf8());
                case InferKind::BinaryDict:
                    return make_dict_converter(binary());
            }
            return turbo::unknown_error("Shouldn't come here");
        }

    protected:
        void SetKind(InferKind kind) {
            kind_ = kind;
            if (kind == InferKind::Binary) {
                // Binary is the catch-all type
                can_loosen_type_ = false;
            }
        }

        InferKind kind_;
        bool can_loosen_type_;
        const ConvertOptions &options_;
    };

}  // namespace nebula::csv
